from math import pow
import numpy as np
from collections import Counter, OrderedDict
import copy


class TreeNode:
    def __init__(self,
                 root=None,
                 attribute=None,
                 property=None,
                 height=1,
                 pos=0):
        self.root = root  #父节点
        self.children = []  #子节点列表
        self.label = None  #节点分类
        self.attribute = attribute  #节点标签
        self.property = property  #对应父节点标签的属性，
        self.height = height  # 树深
        self.pos = pos  #相对位置


def count(a):  #返回是否值相同，数量最多的值
    c = Counter(a)
    if len(c) == 1:
        return True, a[0]
    else:
        return False, c.most_common(1)[0][0]


def split(d, t, a, p):
    #根据标签属性划分数据，并删除划分后的标签
    #用空间换时间
    D, target, _D, _target = [], [], [], []
    for i in range(d.shape[0]):
        if d[i][a] == p:
            _D.append(d[i])
            _target.append(t[i])
        else:
            D.append(d[i])
            target.append(t[i])
    D = np.array(D)
    target = np.array(target)
    #可能出现划分为空
    if len(_D) != 0:
        _D = np.array(_D)
        if len(_D.shape) == 1:
            _D = np.delete(_D, a)
        else:
            _D = np.delete(_D, a, axis=1)
    _target = np.array(_target)
    return D, target, _D, _target


def chooseBestSplit(D, target, A):
    #最佳划分，如果有最佳划分则返回划分类，没有返回则返回None
    #CART算法,Gini指数
    temp = []
    lenD = D.shape[0]
    for i in A:
        Gini_index = 0
        index = list(A).index(i)
        for j in A[i]:
            D, target, _D, _target = split(D, target, index, j)
            if len(_D) != 0:
                t = len(_target)
                sum = 0
                c = Counter(_target)
                for k in c:
                    sum += pow(c[k] / t, 2)
                Gini = 1 - sum
                Gini_index += _D.shape[0] / lenD * Gini
        temp.append(Gini_index)
    if max(temp) == min(temp):
        return None
    else:
        return list(A)[temp.index(min(temp))]


def TreeGenerate(D,
                 target,
                 attribute,
                 node,
                 max_depth,
                 min_leaf_simples=0,
                 h=1):
    A = copy.deepcopy(attribute)
    flag, C = count(target)
    if flag or len(A) == 0 or chooseBestSplit(D, target,
                                              A) == None or h == max_depth:
        node.label = C
        return
    BestSplit = chooseBestSplit(D, target, A)
    node.attribute = BestSplit
    #对于离散的标签，从集合剔除划分标签
    index = list(A).index(BestSplit)
    _A = A.pop(BestSplit)
    for i in range(len(_A)):
        node.children.append(
            TreeNode(root=node, property=_A[i], height=h + 1, pos=i))
        D, target, _D, _target = split(D, target, index, _A[i])
        if len(_target) < min_leaf_simples:
            node.children = []
            node.label = C
            return
        if len(_target) == 0:
            node.children[i].label = C
        else:
            TreeGenerate(_D, _target, A, node.children[i], max_depth,
                         min_leaf_simples, h + 1)


def dataRead(path, usecols, targetCol):
    D = np.loadtxt(path,
                   dtype="str",
                   delimiter=",",
                   skiprows=1,
                   usecols=usecols,
                   encoding='utf-8')
    target = np.loadtxt(path,
                        dtype="str",
                        delimiter=",",
                        skiprows=1,
                        usecols=targetCol,
                        encoding='utf-8')
    #利用有序字典，构建标签和对应属性集合
    rowA = np.loadtxt("table_4.2.csv",
                      dtype="str",
                      delimiter=",",
                      usecols=usecols,
                      encoding='utf-8',
                      max_rows=1)
    A = OrderedDict()
    for i in range(len(rowA)):
        A[rowA[i]] = list(Counter(D[:, i]))
    return D, target, A


D, target, A = dataRead("table_4.2.csv", (1, 2, 3, 4, 5, 6), 9)

D_train = D[:-1, :]
target_train = target[:-1]
D_test = D[-1, :]
target_test = target[-1]

root = TreeNode()
TreeGenerate(D=D,
             target=target,
             attribute=A,
             max_depth=4,
             min_leaf_simples=0,
             node=root)


def predict(d):
    node = root
    while 1:
        if node.label != None:
            return node.label
        index = list(A).index(node.attribute)
        for i in node.children:
            if i.property == d[index]:
                node = i
                break
